#include "StateEstimator.h"
#include <string>
#include <cmath>
#include <Eigen/Dense>
#include <iostream>

using namespace std;
using namespace Eigen;

StateEstimator::StateEstimator() {
  double t = DELTA_T_OTHERS;
  // 初始化
  IMUsetflag_ = false;
  Encoderflag_ = false;
  Calculatedflag_ = false;

  R_B_O_.setIdentity(3, 3);
  R_B_imu_.setIdentity(3, 3);
  R_imu0_O_.setIdentity(3, 3);
  R_t0_imu_imu0_.setIdentity(3, 3);
  p_i_B_.setZero(3, 1);
  p_dot_i_B_.setZero(3, 1);
  x_k_.setZero(28, 1);
  a_com_B_.setZero(3, 1);
  a_com_O_.setZero(3, 1);
  omega_OB_B_.setZero(3, 1);
  omega_OB_O_.setZero(3, 1);
  omega_OBX_B_.setZero(3, 3);
  // BUG: 未初始化导致getp_stend_O()返回错误值,请根据真实数据进行初始化
  p_stend_O_last_.setZero(4, 3);

  // Kalman filter init
  A_.setIdentity();
  A_.block<3, 3>(0, 3) = t * Matrix3d::Identity();

  B_.setIdentity();
  B_.block<3, 3>(3, 0) = t * Matrix3d::Identity();

  H_.block<3, 3>(0, 0).setIdentity(3, 3);
  H_.block<3, 3>(3, 0).setIdentity(3, 3);
  H_.block<3, 3>(6, 0).setIdentity(3, 3);
  H_.block<3, 3>(9, 0).setIdentity(3, 3);
  H_.block<3, 3>(0, 3).setZero();
  H_.block<3, 3>(3, 3).setZero();
  H_.block<3, 3>(6, 3).setZero();
  H_.block<3, 3>(9, 3).setZero();

  H_.block<12, 12>(0, 6).setIdentity(12, 12);
  H_.block<12, 12>(0, 6) = -H_.block<12, 12>(0, 6);
  // H_(12-23)
  H_.block<3, 3>(12, 0).setZero();
  H_.block<3, 3>(15, 0).setZero();
  H_.block<3, 3>(18, 0).setZero();
  H_.block<3, 3>(21, 0).setZero();
  H_.block<3, 3>(12, 3).setIdentity(3, 3);
  H_.block<3, 3>(15, 3).setIdentity(3, 3);
  H_.block<3, 3>(18, 3).setIdentity(3, 3);
  H_.block<3, 3>(21, 3).setIdentity(3, 3);

  H_.block<12, 12>(12, 6).setZero();
  // H_(24-27)
  H_.block<4, 18>(24, 0).setZero();
  H_(24, 8) = 1;
  H_(25, 11) = 1;
  H_(26, 14) = 1;
  H_(27, 17) = 1;

  Q_.setIdentity(); // XXX: 根据传感器误差数据调参

  R_ = 10 * Matrix<double, 28, 28>::Identity(); // XXX: 根据传感器误差数据调参

  P_ = 10 * Matrix<double, 18, 18>::Identity(); // XXX: 调参

  K_.setZero();

  u_ << 0, 0, -9.8;
}

StateEstimator::~StateEstimator() {}

void StateEstimator::setIMUdata(const Vector3d &Eulerangles,
                                const Vector3d &a_original,
                                const Vector3d &omega_original) {
  double phi = Eulerangles(0);
  double theta = Eulerangles(1);
  double psi = Eulerangles(2);

  R_imu_imu0_ << cos(theta) * cos(psi),
      sin(phi) * sin(theta) * cos(psi) - cos(phi) * sin(psi),
      cos(phi) * sin(theta) * cos(psi) + sin(phi) * sin(psi),
      cos(theta) * sin(psi),
      sin(phi) * sin(theta) * sin(psi) + cos(phi) * cos(psi),
      cos(phi) * sin(theta) * sin(psi) - sin(phi) * cos(psi), -sin(theta),
      sin(phi) * cos(theta), cos(phi) * cos(theta);

  a_original_ = a_original;

  omega_original_ = omega_original;

  IMUsetflag_ = true;
}

void StateEstimator::setEncoderdata(const Matrix<double, 4, 3> &q_j_i,
                                    const Matrix<double, 4, 3> &q_dot_j_i) {
  q_j_i_ = q_j_i;
  q_dot_j_i_ = q_dot_j_i;

  Encoderflag_ = true;
}

void StateEstimator::calculateAll() {
  if (IMUsetflag_ != true || Encoderflag_ != true) {
    cerr << "StateEstimator: set necessary data before calculate!" << endl;
    exit(1);
  }
  // R_B_imu_ = R("x", 180.0); // 原论文定义
  // double psi_init = -asin(R_t0_imu_imu0_(1, 0)); // asin返回弧度
  // psi_init = psi_init * 180.0 / 3.1415326;
  // R_imu0_O_ = (R("z", psi_init)) * (R("x", 180)); // 原论文定义

  // R_B_imu_.setIdentity(3, 3); // spot定义
  // R_imu0_O_.setIdentity(3, 3); // spot定义
  // 计算旋转矩阵
  // R_B_O_ = R_imu0_O_ * R_imu_imu0_ * R_B_imu_;
  R_B_O_ = R_imu_imu0_; // ATTENTION:

  R_B_O_.setIdentity(3, 3); // BUG
  // 与原论文定义不一致，可能存在与后续变量不匹配的bug

  // Matrix3d R_imu_B_;
  // R_B_imu_ = R("x", -180.0);

  // a_com_B_ = R_imu_B_ * a_original;
  a_com_B_ = a_original_; // spot定义

  // Matrix3d R_imu_B_;
  // R_imu_B_ = R("x", -180.0);

  // omega_OB_B_ = R_imu_B_ * omega_original;
  omega_OB_B_ = omega_original_; // spot定义

  a_com_O_ = R_B_O_ * a_com_B_;
  omega_OB_O_ = R_B_O_ * omega_OB_B_;

  // clang-format off
    omega_OBX_B_ << 0, -omega_OB_B_(2), omega_OB_B_(1),
                    omega_OB_B_(2), 0, -omega_OB_B_(0),
                    -omega_OB_B_(1), omega_OB_B_(0), 0;
  // clang-format on

  // z_k_(0-11)
  z_k_.segment(0, 3) = -R_B_O_ * getp_i_B(0, q_j_i_.row(0));
  z_k_.segment(3, 3) = -R_B_O_ * getp_i_B(1, q_j_i_.row(1));
  z_k_.segment(6, 3) = -R_B_O_ * getp_i_B(2, q_j_i_.row(2));
  z_k_.segment(9, 3) = -R_B_O_ * getp_i_B(3, q_j_i_.row(3));

  // z_k_(12-23)
  z_k_.segment(12, 3) =
      -R_B_O_ * (getomega_OBX_B() * getp_i_B(0, q_j_i_.row(0))) +
      getp_dot_i_B(0, q_j_i_.row(0), q_dot_j_i_.row(0));
  z_k_.segment(15, 3) =
      -R_B_O_ * (getomega_OBX_B() * getp_i_B(1, q_j_i_.row(1))) +
      getp_dot_i_B(1, q_j_i_.row(1), q_dot_j_i_.row(1));
  z_k_.segment(18, 3) =
      -R_B_O_ * (getomega_OBX_B() * getp_i_B(2, q_j_i_.row(2))) +
      getp_dot_i_B(2, q_j_i_.row(2), q_dot_j_i_.row(2));
  z_k_.segment(21, 3) =
      -R_B_O_ * (getomega_OBX_B() * getp_i_B(3, q_j_i_.row(3))) +
      getp_dot_i_B(3, q_j_i_.row(3), q_dot_j_i_.row(3));

  // z_k_(24-27)
  for (int i = 0; i < 4; i++) {
    z_k_(i + 24) = 0;
  }

  Calculatedflag_ = true;
}

Matrix3d StateEstimator::getR_B_O() { return R_B_O_; }

Vector3d StateEstimator::geta_com_B() { return a_com_B_; }

Vector3d StateEstimator::getomega_OB_B() { return omega_OB_B_; }

/**
 * Calculates the position vector p_i_B for a given index i and a given vector
 * q_j_i.
 *
 * @param i The index value between 0 ~ 3.
 * @param q_j_i The motor angles of the leg i: (gamma, alfa, beta).
 * @return The position vector p_i_B.
 */
Vector3d StateEstimator::getp_i_B(int i, const Vector3d &q_j_i) {
  // int ksi = (i == 2 || i == 4) ? 1 : -1;
  // int del = (i == 1 || i == 2) ? 1 : -1;
  //(2.48)
  // p_i_B_(0) = -L_2 * sin(q_j_i(1)) - L_3 * sin(q_j_i(1) + q_j_i(2)) + del *
  // H_X;
  // p_i_B_(1) = ksi * L_1 * cos(q_j_i(0)) +
  //             L_3 * sin(q_j_i(0)) * cos(q_j_i(1) + q_j_i(2)) +
  //             L_2 * cos(q_j_i(1)) * sin(q_j_i(0)) + ksi * H_Y;
  // p_i_B_(2) = ksi * L_1 * sin(q_j_i(0)) -
  //             L_3 * cos(q_j_i(0)) * cos(q_j_i(1) + q_j_i(2)) -
  //             L_2 * cos(q_j_i(0)) * cos(q_j_i(1));

  p_i_B_ = legToBody(i, forwardKinematics(i, q_j_i));

  // cout << "leg " << i << " p_i_B_ " << p_i_B_.transpose() << endl;
  return p_i_B_;
}

/**
 * Calculates the foot velocity vector of leg i in the body system.
 *
 * @param i The index of the leg between 0 ~ 3.
 * @param q_j_i The position vector of point i relative to point j.
 * @param q_dot_j_i The derivative of the position vector q_j_i.
 * @return The derivative of the position vector of point i in the body frame B.
 */
Vector3d StateEstimator::getp_dot_i_B(int i, const Vector3d &q_j_i,
                                      const Vector3d &q_dot_j_i) {
  int ksi = (i == 1 || i == 3) ? 1 : -1;
  // // int del = (i == 1 || i == 2) ? 1 : -1;  // not used in this function
  // //(2.49)
  // Matrix3d tmpmat;
  // tmpmat(0, 0) = 0;
  // tmpmat(0, 1) = -L_2 * cos(q_j_i(1)) - L_3 * cos(q_j_i(1) + q_j_i(2));
  // tmpmat(0, 2) = -L_3 * cos(q_j_i(1) + q_j_i(2));
  // tmpmat(1, 0) = ksi * L_1 * sin(q_j_i(0)) +
  //                L_3 * cos(q_j_i(0)) * cos(q_j_i(1) + q_j_i(2)) +
  //                L_2 * cos(q_j_i(1)) * cos(q_j_i(0));
  // tmpmat(1, 1) = -L_3 * sin(q_j_i(0)) * sin(q_j_i(1) + q_j_i(2)) -
  //                L_2 * sin(q_j_i(1)) * sin(q_j_i(0));
  // tmpmat(1, 2) = -L_3 * sin(q_j_i(0)) * sin(q_j_i(1) + q_j_i(2));
  // tmpmat(2, 0) = ksi * L_1 * cos(q_j_i(0)) +
  //                L_3 * sin(q_j_i(0)) * cos(q_j_i(1) + q_j_i(2)) +
  //                L_2 * cos(q_j_i(1)) * sin(q_j_i(0));
  // tmpmat(2, 1) = L_3 * cos(q_j_i(0)) * sin(q_j_i(1) + q_j_i(2)) +
  //                L_2 * sin(q_j_i(1)) * cos(q_j_i(0));
  // tmpmat(2, 2) = L_3 * cos(q_j_i(0)) * sin(q_j_i(1) + q_j_i(2));

  // p_dot_i_B_ = tmpmat * q_dot_j_i;

  // p_dot_i_B_ = calculateJacobian(i, q_j_i) * q_dot_j_i;

  double x = -L_2 * sin(q_j_i(1)) * q_dot_j_i(1) +
             L_3 * sin(q_j_i(1) + q_j_i(2)) * (q_dot_j_i(1) + q_dot_j_i(2));

  double y =
      ksi * (L_2 * cos(q_j_i(1)) * q_dot_j_i(1) -
             L_3 * cos(q_j_i(1) + q_j_i(2)) * (q_dot_j_i(1) + q_dot_j_i(2))) *
          cos(q_j_i(0)) -
      ksi * (L_2 * sin(q_j_i(1)) - L_3 * sin(q_j_i(1) + q_j_i(2))) *
          sin(q_j_i(0) * q_dot_j_i(0)) +
      ksi * L_1 * cos(q_j_i(0)) * q_dot_j_i(0);

  double z = -(L_2 * cos(q_j_i(1)) * q_dot_j_i(1) -
               L_3 * cos(q_j_i(1) + q_j_i(2)) * (q_dot_j_i(1) + q_dot_j_i(2))) *
                 sin(q_j_i(0)) -
             (L_2 * sin(q_j_i(1)) - L_3 * sin(q_j_i(1) + q_j_i(2))) *
                 cos(q_j_i(0)) * q_dot_j_i(0) -
             L_1 * sin(q_j_i(0)) * q_dot_j_i(0);

  p_dot_i_B_ << x, y, z;

  cout << "leg " << i << " p_dot_i_B_ " << p_dot_i_B_.transpose() << endl;

  return p_dot_i_B_;
}

Vector3d StateEstimator::geta_com_O() { return a_com_O_; }

Vector3d StateEstimator::getomega_OB_O() { return omega_OB_O_; }

Matrix3d StateEstimator::getomega_OBX_B() { return omega_OBX_B_; }

void StateEstimator::kalman_update(int times) {
  if (Calculatedflag_ != true) {
    cerr << "StateEstimator: calculateAll() before kalman_update()!" << endl;
    exit(1);
  }
  //  腿编号
  //  1   0
  //  3   2

  // Kalman filter
  for (int i = 0; i < times; i++) {
    x_k_ = A_ * x_k_ + B_ * u_;
    P_ = A_ * P_ * A_.transpose() + Q_;

    K_ = (P_ * H_.transpose()) * (H_ * P_ * H_.transpose() + R_).inverse();
    x_k_ = x_k_ + K_ * (z_k_ - H_ * x_k_);
    P_ = (Matrix<double, 18, 18>::Identity() - K_ * H_) * P_;
  }
}

Matrix<double, 4, 3> StateEstimator::getp_stend_O(const Vector4d &s_fai) {
  Matrix<double, 4, 3> p_stend_O;
  for (int i = 0; i < 4; i++) {
    if (s_fai(i) == 1) {
      p_stend_O.row(i) = R_B_O_ * getp_i_B(i, q_j_i_.row(i));
      p_stend_O_last_.row(i) = p_stend_O.row(i);
    } else {
      p_stend_O.row(i) = p_stend_O_last_.row(i);
    }
  }
  return p_stend_O;
}

Matrix<double, 18, 1> StateEstimator::getx_k() { return x_k_; }

Matrix<double, 4, 3> StateEstimator::getp_foot() {
  Matrix<double, 4, 3> p_foot;
  for (int i = 0; i < 4; i++) {
    p_foot.row(i) = x_k_.segment(6 + i * 3, 3).transpose();
  }
  return p_foot;
}
